import numpy as np
import torch
from torch.utils.data import Dataset

from model.ours_LFDS.utils import get_lds_kernel_window
from scipy.ndimage import convolve1d

class MyDataset(Dataset):
    def __init__(self, X, y=None, reweight='none', lds=False, lds_kernel='gaussian', lds_ks=5, lds_sigma=2):
        self.X = X
        self.y = y
        self.lds=lds
        self.weights=None

        if lds:
            self.weights = self._prepare_weights(reweight=reweight, lds=lds, lds_kernel=lds_kernel, lds_ks=lds_ks,
                                             lds_sigma=lds_sigma)
        print("")

    def __len__(self):
        return len(self.X)

    def __getitem__(self, index):
        # print(self.X[index])
        # print(self.y.iloc[index])

        # weight = np.asarray([self.weights[index]]).astype('float32') if self.weights is not None else np.asarray(
        #     [np.float32(1.)])
        weight = np.asarray(self.weights[index]).astype('float32') if self.weights is not None else np.asarray(
            np.float32(1.))
        # weight=np.asarray(self.weights)

        return self.X[index], self.y[index], weight

        # if self.lds:
        #     return self.X[index], self.y[index], weight
        # else:
        #     return self.X[index], self.y[index]

    def _prepare_weights(self, reweight, max_target=1005, lds=False, lds_kernel='gaussian', lds_ks=5, lds_sigma=2):
        assert reweight in {'none', 'inverse', 'sqrt_inv'}
        assert reweight != 'none' if lds else True, \
            "Set reweight to \'sqrt_inv\' (default) or \'inverse\' when using LDS"

        # value_dict = {x: 0 for x in range(max_target)}
        value_dict = {x: 0 for x in range(max_target)}
        labels = self.y
        for label in labels:
            value_dict[min(max_target - 1, int(label*1000))] += 1
            # if str(label) not in value_dict:
            #     value_dict[str(label)] = 1
            # else:
            #     value_dict[str(label)] += 1
        if reweight == 'sqrt_inv':
            value_dict = {k: np.sqrt(v) for k, v in value_dict.items()}
        elif reweight == 'inverse':
            value_dict = {k: np.clip(v, 0, 1000) for k, v in value_dict.items()}  # clip weights for inverse re-weight
        num_per_label = [value_dict[min(max_target - 1, int(label*1000))] for label in labels]
        # num_per_label = [value_dict[str(label)] for label in labels]
        if not len(num_per_label) or reweight == 'none':
            return None
        print(f"Using re-weighting: [{reweight.upper()}]")

        if lds:
            lds_kernel_window = get_lds_kernel_window(lds_kernel, lds_ks, lds_sigma)
            print(f'Using LDS: [{lds_kernel.upper()}] ({lds_ks}/{lds_sigma})')
            smoothed_value = convolve1d(
                np.asarray([v for _, v in value_dict.items()]), weights=lds_kernel_window, mode='constant')
            num_per_label = [smoothed_value[min(max_target - 1, int(label*1000))] for label in labels]
            # num_per_label = [smoothed_value[str(label)] for label in labels]

        weights = [np.float32(1 / x) for x in num_per_label]
        scaling = len(weights) / np.sum(weights)
        weights = [scaling * x for x in weights]
        print('weights: ',weights)
        return weights